/*
 * Copyright (c) 2023-2024 elsfs Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.elsfs.cloud.screw.query.db2;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import javax.sql.DataSource;
import org.apache.commons.lang3.StringUtils;
import org.elsfs.cloud.common.util.exception.QueryException;
import org.elsfs.cloud.common.util.lang.Assert;
import org.elsfs.cloud.common.util.lang.ExceptionUtils;
import org.elsfs.cloud.common.util.sql.JdbcUtils;
import org.elsfs.cloud.screw.constant.ScrewConstants;
import org.elsfs.cloud.screw.mapping.Mapping;
import org.elsfs.cloud.screw.metadata.Column;
import org.elsfs.cloud.screw.metadata.Database;
import org.elsfs.cloud.screw.metadata.PrimaryKey;
import org.elsfs.cloud.screw.metadata.Table;
import org.elsfs.cloud.screw.query.AbstractDatabaseQuery;
import org.elsfs.cloud.screw.query.db2.model.Db2ColumnModel;
import org.elsfs.cloud.screw.query.db2.model.Db2DatabaseModel;
import org.elsfs.cloud.screw.query.db2.model.Db2PrimaryKeyModel;
import org.elsfs.cloud.screw.query.db2.model.Db2TableModel;

/**
 * db2 数据库查询
 *
 * @author zeng
 */
public class Db2DataBaseQuery extends AbstractDatabaseQuery {
  /**
   * 构造函数
   *
   * @param dataSource {@link DataSource}
   */
  public Db2DataBaseQuery(DataSource dataSource) {
    super(dataSource);
  }

  /**
   * 获取数据库
   *
   * @return {@link Database} 数据库信息
   */
  @Override
  public Database getDataBase() throws QueryException {
    Db2DatabaseModel model = new Db2DatabaseModel();
    model.setDatabase(StringUtils.trim(this.getSchema()));
    return model;
  }

  /**
   * 获取表信息
   *
   * @return {@link List} 所有表信息
   */
  @Override
  public List<Db2TableModel> getTables() {
    ResultSet resultSet = null;
    try {
      // 查询
      resultSet =
          getMetaData()
              .getTables(
                  getCatalog(), getSchema(), ScrewConstants.PERCENT_SIGN, new String[] {"TABLE"});
      // 映射
      return Mapping.convertList(resultSet, Db2TableModel.class);
    } catch (SQLException e) {
      throw ExceptionUtils.mpe(e);
    } finally {
      JdbcUtils.close(resultSet, connection);
    }
  }

  /**
   * 获取列信息
   *
   * @param table {@link String} 表名
   * @return {@link List} 表字段信息
   * @throws QueryException QueryException
   */
  @Override
  public List<Db2ColumnModel> getTableColumns(String table) throws QueryException {
    Assert.notEmpty(table, "Table name can not be empty!");
    ResultSet resultSet = null;
    try {
      // 查询所有
      List<Db2ColumnModel> list;
      if (ScrewConstants.PERCENT_SIGN.equals(table)) {
        List<String> tableNames =
            getTables().stream().map(Table::getTableName).toList().stream().distinct().toList();
        // db2 getMetaData().getColumns 中 无法获取表名,循环单张表获取
        list = new ArrayList<>();
        for (String tableName : tableNames) {
          list.addAll(getTableColumns(tableName));
        }
      } else {
        if (!this.columnsCaching.containsKey(table)) {
          // 查询
          resultSet =
              getMetaData()
                  .getColumns(getCatalog(), getSchema(), table, ScrewConstants.PERCENT_SIGN);
          // 映射
          list = Mapping.convertList(resultSet, Db2ColumnModel.class);
          // 单表查询
          String sql =
              "SELECT COLNAME as NAME,TABNAME as TABLE_NAME,REMARKS,LENGTH as COLUMN_LENGTH,"
                  + " TYPENAME ||'('|| LENGTH ||')'  as COLUMN_TYPE,SCALE as DECIMAL_DIGITS   FROM"
                  + " SYSCAT.COLUMNS WHERE TABSCHEMA='"
                  + StringUtils.trim(getSchema())
                  + "' and TABNAME = '%s' ORDER BY COLNO";
          resultSet = prepareStatement(String.format(sql, table)).executeQuery();
          // db2 ColumnName 获取的不是 as column 值,使用ColumnLabel获取
          List<Db2ColumnModel> inquires =
              Mapping.convertListByColumnLabel(resultSet, Db2ColumnModel.class);
          for (Db2ColumnModel i : list) {
            inquires.forEach(
                j -> {
                  // 列名一致
                  if (i.getColumnName().equals(j.getColumnName())) {
                    // 放入备注
                    i.setRemarks(j.getRemarks());
                    i.setColumnLength(j.getColumnLength());
                    i.setColumnType(j.getColumnType());
                    i.setTableName(j.getTableName());
                    i.setDecimalDigits(j.getDecimalDigits());
                    if ("NO".equals(i.getNullable())) {
                      i.setNullable("0");
                    }
                  }
                });
          }
          this.columnsCaching.put(
              table,
              list.stream()
                  .filter(i -> StringUtils.isNotBlank(i.getColumnName()))
                  .collect(Collectors.toList()));
        } else {
          list =
              this.columnsCaching.get(table).stream()
                  .map(c -> (Db2ColumnModel) c)
                  .collect(Collectors.toList());
        }
      }
      return list;
    } catch (SQLException e) {
      throw ExceptionUtils.mpe(e);
    } finally {
      JdbcUtils.close(resultSet, this.connection);
    }
  }

  /**
   * 获取所有列信息
   *
   * @return {@link List} 表字段信息
   * @throws QueryException QueryException
   */
  @Override
  public List<? extends Column> getTableColumns() throws QueryException {
    return getTableColumns(ScrewConstants.PERCENT_SIGN);
  }

  /**
   * 根据表名获取主键
   *
   * @param table {@link String}
   * @return {@link List}
   * @throws QueryException QueryException
   */
  @Override
  public List<? extends PrimaryKey> getPrimaryKeys(String table) throws QueryException {
    ResultSet resultSet = null;
    try {
      // 查询
      resultSet = getMetaData().getPrimaryKeys(getCatalog(), getSchema(), table);
      // 映射
      return Mapping.convertList(resultSet, Db2PrimaryKeyModel.class);
    } catch (SQLException e) {
      throw ExceptionUtils.mpe(e);
    } finally {
      JdbcUtils.close(resultSet, this.connection);
    }
  }

  /**
   * 根据表名获取主键
   *
   * @return {@link List}
   * @throws QueryException QueryException
   */
  @Override
  public List<? extends PrimaryKey> getPrimaryKeys() throws QueryException {
    ResultSet resultSet = null;
    try {
      // 由于单条循环查询存在性能问题，所以这里通过自定义SQL查询数据库主键信息
      String sql =
          "SELECT TABSCHEMA as TABLE_SCHEM, TABNAME as TABLE_NAME ,COLNAME as COLUMN_NAME,KEYSEQ as"
              + " KEY_SEQ   FROM SYSCAT.COLUMNS WHERE TABSCHEMA = '%s' AND KEYSEQ IS NOT NULL ORDER"
              + " BY KEYSEQ";
      resultSet = prepareStatement(String.format(sql, getDataBase().getDatabase())).executeQuery();
      return Mapping.convertListByColumnLabel(resultSet, Db2PrimaryKeyModel.class);
    } catch (SQLException e) {
      throw new QueryException(e);
    } finally {
      JdbcUtils.close(resultSet);
    }
  }
}
